`include "Header.svh"

/***********************************************************
Moodule: PE
Function: 8 multiple and add tree
************************************************************/
module PE(
  input logic clk,
  input logic rst_n,
  input logic data_in_valid,//Ctrol
  input logic [1:0] Cal_Mode,//Ctrol
  input logic [4:0] K_idx,//Ctrol
  input pack_data a[7:0],
  input pack_data b[7:0],
  input logic psum_en,
  input idx_t position_di,
  output idx_t position_do,
  output pack_data pe_result,
  output logic pe_valid
  );
  
  /******************************************************************************************************
                    ------------------------Ctrol Unit------------------------
  *******************************************************************************************************/
  // Mode selection (judgment)
    // `INT4_MODE = 'd0,
    // `INT8_MODE = 'd1,
    // `FP16_MODE = 'd2,
    // `FP32_MODE = 'd3;
  //Behavioral judgment
    //add_valid.The 8 numbers need to undergo a total of 6 additions to obtain the result.
        //8 to 4 
        //4 to 2 
        //2 to 1 
        //accumulation 
        //Bias
        logic intadd_valid;
    //mul_valid.The 8 numbers need to undergo a total of 2 multiplications to obtain the result.
        //8 to 8
        logic intmul_valid;
    always_comb begin : Mode_judgment
      if(Cal_Mode == `INT4_MODE || Cal_Mode == `INT8_MODE)begin
        intadd_valid = 1'b1;
        intmul_valid = 1'b1;
      end else begin
        intadd_valid = 1'b0;
        intmul_valid = 1'b0;
      end
    end
  //Flag
    //last_block.final block flag.Judge the accumulation
      logic accumulation_enable;
      assign accumulation_enable = ~(K_idx == 1);
      logic last_block;
      assign last_block = (position_reg[4].k1_idx == (K_idx-1));
    //The next state judgment condition
  //Pipe Line //One-hot encoding
    logic psum_en_reg [5:0];
    idx_t position_reg[5:0];
    logic valid_reg   [5:0];  
    always_ff @( posedge clk ) begin : pipe_line
      if(!rst_n)begin
        for(integer i =0;i<6;i=i+1)begin
          psum_en_reg[i] <= 'd0;
          position_reg[i] <= 'd0;
          valid_reg[i] <= 'd0;
        end
      end else begin
        if(data_in_valid)begin
          psum_en_reg[0] <= psum_en;
          position_reg[0] <= position_di;
          valid_reg[0] <= data_in_valid;
        end else begin
          psum_en_reg[0] <= 'd0;
          position_reg[0] <= 'd0;
          valid_reg[0] <= 'd0;
        end
        for(integer i=1;i<6;i=i+1)begin
          psum_en_reg[i] <= psum_en_reg[i-1];
          position_reg[i] <= position_reg[i-1];
          valid_reg[i] <= valid_reg[i-1];
        end
      end
    end
  /******************************************************************************************************
                    ------------------------Data Path------------------------
  *******************************************************************************************************/
  // instantiation
    //Int4 with Int8 Use the same signed 16
    logic signed [15:0] product[`K0-1:0];
    generate
      for(genvar p = 0;p<`K0;p=p+1)begin
        int_mul INT_mul_inst(
          .clk    (clk),  
          .rst_n  (rst_n),
          .calc_en(data_in_valid && intadd_valid)  , // 有符号位取的数据是这样的，没得办法的说
          .a      ({a[p].data.bit16[1][15],a[p].data.bit16[0][14:0]}),
          .b      ({b[p].data.bit16[1][15],b[p].data.bit16[0][14:0]}),
          .c      (product[p]),
          .done   ()
        );
      end
    endgenerate
    //addtree1 Stage1 8 to 4
    logic signed [15:0] int_add_stage1[4];
    generate
      for(genvar p = 0;p < 4;p=p+1)begin
        int33_addition INT_add1_inst(
          .clk    (clk) ,
          .rst_n  (rst_n) ,
          .calc_en(valid_reg[0] && intadd_valid) ,
          .a      (product[p]) ,
          .b      (product[p+4]) ,
          .c      (int_add_stage1[p]) ,
          .done   ()  
        );
      end
    endgenerate
    //addtree2 Stage2 4 to 2
    logic signed [15:0]  int_add_stage2[2];
    generate
      for(genvar p=0; p<2;p=p+1)begin
        int33_addition INT_add2_inst(
          .clk    (clk) ,
          .rst_n  (rst_n) ,
          .calc_en(valid_reg[1] && intadd_valid) ,
          .a      (int_add_stage1[p]) ,
          .b      (int_add_stage1[p+2]) ,
          .c      (int_add_stage2[p]) ,
          .done   ()  
        );
      end
    endgenerate
    //addtree3 Stage3 2 to 1
    logic signed [15:0]  int_add_stage3;
    int33_addition INT_add3_inst(
      .clk    (clk) ,
      .rst_n  (rst_n) ,
      .calc_en(valid_reg[2] && intadd_valid) ,
      .a      (int_add_stage2[0]) ,
      .b      (int_add_stage2[1]) ,
      .c      (int_add_stage3) ,
      .done   ()  
    );
    //Accumulation psum
    logic signed [15:0] int_add_stage3_1b;
    always_ff @( posedge clk ) begin : delay_1b
      if(!rst_n)begin
        int_add_stage3_1b <= 'd0;
      end else begin
        int_add_stage3_1b <= int_add_stage3;
      end
    end
    logic signed [15:0]  int_psum;
    int33_addition INT_psum_inst(
      .clk    (clk) ,
      .rst_n  (rst_n) ,
      .calc_en(valid_reg[3] && intadd_valid) ,
      .a      (int_add_stage3) ,
      .b      ((psum_en_reg[3] && ~(position_reg[3].k1_idx == 'd0))? int_psum : 'd0) ,
      .c      (int_psum) ,
      .done   ()  
    );
  // Result selection and processing
    //pe_result pe_valid
    always_ff @( posedge clk ) begin : Result_process
      if(valid_reg[4] && last_block)begin
        //Final Result done
        if(Cal_Mode == `INT8_MODE || Cal_Mode == `INT4_MODE)begin
          pe_result.data.bit16[0] <= int_psum;
          pe_valid <= 1'b1;
        end
      end else begin
        pe_result <= 'd0;
        pe_valid <= 'd0;
      end
    end
    assign position_do = position_reg[5];
endmodule